/* -*- Mode:C++; c-file-style:"gnu"; indent-tabs-mode:nil; -*- */
/*
 * Copyright (c) 2006, 2009 INRIA
 * Copyright (c) 2009 MIRKO BANCHI
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation;
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 * Author: Mathieu Lacage <mathieu.lacage@sophia.inria.fr>
 * Author: Mirko Banchi <mk.banchi@gmail.com>
 */
#include "StaCtrlWifiMac.h"

#include "ns3/log.h"
#include "ns3/simulator.h"
#include "ns3/string.h"
#include "ns3/pointer.h"
#include "ns3/boolean.h"
#include "ns3/trace-source-accessor.h"

#include "qos-tag.h"
#include "mac-low.h"
#include "dcf-manager.h"
#include "mac-rx-middle.h"
#include "mac-tx-middle.h"
#include "wifi-mac-header.h"
#include "msdu-aggregator.h"
#include "amsdu-subframe-header.h"
#include "mgt-headers.h"

NS_LOG_COMPONENT_DEFINE("StaCtrlWifiMac");

/*
 * The state machine for this STA is:
 --------------                                          -----------
 | Associated |   <--------------------      ------->    | Refused |
 --------------                        \    /            -----------
 \                                   \  /
 \    -----------------     -----------------------------
 \-> | Beacon Missed | --> | Wait Association Response |
 -----------------     -----------------------------
 \                       ^
 \                      |
 \    -----------------------
 \-> | Wait Probe Response |
 -----------------------
 */

namespace ns3
{

NS_OBJECT_ENSURE_REGISTERED(StaCtrlWifiMac);

TypeId StaCtrlWifiMac::GetTypeId(void)
{
	static TypeId tid =
			TypeId("ns3::StaCtrlWifiMac").SetParent<RegularCtrlWifiMac>().AddConstructor<
					StaCtrlWifiMac>().AddAttribute("ProbeRequestTimeout",
					"The interval between two consecutive probe request attempts.",
					TimeValue(Seconds(0.05)),
					MakeTimeAccessor(&StaCtrlWifiMac::m_probeRequestTimeout),
					MakeTimeChecker()).AddAttribute("AssocRequestTimeout",
					"The interval between two consecutive assoc request attempts.",
					TimeValue(Seconds(0.5)),
					MakeTimeAccessor(&StaCtrlWifiMac::m_assocRequestTimeout),
					MakeTimeChecker()).AddAttribute("MaxMissedBeacons",
					"Number of beacons which much be consecutively missed before "
							"we attempt to restart association.",
					UintegerValue(10),
					MakeUintegerAccessor(&StaCtrlWifiMac::m_maxMissedBeacons),
					MakeUintegerChecker<uint32_t>()).AddAttribute(
					"ActiveProbing",
					"If true, we send probe requests. If false, we don't. NOTE: if more than one STA in your simulation is using active probing, you should enable it at a different simulation time for each STA, otherwise all the STAs will start sending probes at the same time resulting in collisions. See bug 1060 for more info.",
					BooleanValue(false),
					MakeBooleanAccessor(&StaCtrlWifiMac::SetActiveProbing),
					MakeBooleanChecker()).AddTraceSource("Assoc",
					"Associated with an access point.",
					MakeTraceSourceAccessor(&StaCtrlWifiMac::m_assocLogger)).AddTraceSource(
					"DeAssoc", "Association with an access point lost.",
					MakeTraceSourceAccessor(&StaCtrlWifiMac::m_deAssocLogger));
	return tid;
}

StaCtrlWifiMac::StaCtrlWifiMac() :
		m_state(BEACON_MISSED), m_probeRequestEvent(), m_assocRequestEvent(), m_beaconWatchdogEnd(
				Seconds(0.0))
{
	NS_LOG_FUNCTION (this);

	// Let the lower layers know that we are acting as a non-AP STA in
	// an infrastructure BSS.
	SetTypeOfStation(STA);

}

StaCtrlWifiMac::~StaCtrlWifiMac()
{
	NS_LOG_FUNCTION (this);
}

void StaCtrlWifiMac::SetMaxMissedBeacons(uint32_t missed)
{
	NS_LOG_FUNCTION (this << missed);
	m_maxMissedBeacons = missed;
}

void StaCtrlWifiMac::SetProbeRequestTimeout(Time timeout)
{
	NS_LOG_FUNCTION (this << timeout);
	m_probeRequestTimeout = timeout;
}

void StaCtrlWifiMac::SetAssocRequestTimeout(Time timeout)
{
	NS_LOG_FUNCTION (this << timeout);
	m_assocRequestTimeout = timeout;
}

void StaCtrlWifiMac::StartActiveAssociation(void)
{
	NS_LOG_FUNCTION (this);
	TryToEnsureAssociated();
}

void StaCtrlWifiMac::SetActiveProbing(bool enable)
{
	NS_LOG_FUNCTION (this << enable);
	if (enable)
	{
		Simulator::ScheduleNow(&StaCtrlWifiMac::TryToEnsureAssociated, this);
	}
	else
	{
		m_probeRequestEvent.Cancel();
	}
}

void StaCtrlWifiMac::SendProbeRequest(void)
{
	NS_LOG_FUNCTION (this);
	WifiMacHeader hdr;
	hdr.SetProbeReq();
	hdr.SetAddr1(Mac48Address::GetBroadcast());
	hdr.SetAddr2(GetAddress());
	hdr.SetAddr3(Mac48Address::GetBroadcast());
	hdr.SetDsNotFrom();
	hdr.SetDsNotTo();
	Ptr<Packet> packet = Create<Packet>();
	MgtProbeRequestHeader probe;
	probe.SetSsid(GetSsid());
	probe.SetSupportedRates(GetSupportedRates());
	packet->AddHeader(probe);

	// The standard is not clear on the correct queue for management
	// frames if we are a QoS AP. The approach taken here is to always
	// use the DCF for these regardless of whether we have a QoS
	// association or not.
	m_dca->Queue(packet, hdr);

	m_probeRequestEvent = Simulator::Schedule(m_probeRequestTimeout,
			&StaCtrlWifiMac::ProbeRequestTimeout, this);
}

void StaCtrlWifiMac::SendAssociationRequest(void)
{
	NS_LOG_FUNCTION (this << GetBssid ());
	WifiMacHeader hdr;
	hdr.SetAssocReq();
	hdr.SetAddr1(GetBssid());
	hdr.SetAddr2(GetAddress());
	hdr.SetAddr3(GetBssid());
	hdr.SetDsNotFrom();
	hdr.SetDsNotTo();
	Ptr<Packet> packet = Create<Packet>();
	MgtAssocRequestHeader assoc;
	assoc.SetSsid(GetSsid());
	assoc.SetSupportedRates(GetSupportedRates());
	packet->AddHeader(assoc);

	// The standard is not clear on the correct queue for management
	// frames if we are a QoS AP. The approach taken here is to always
	// use the DCF for these regardless of whether we have a QoS
	// association or not.
	m_dca->Queue(packet, hdr);

	m_assocRequestEvent = Simulator::Schedule(m_assocRequestTimeout,
			&StaCtrlWifiMac::AssocRequestTimeout, this);
}

void StaCtrlWifiMac::TryToEnsureAssociated(void)
{
	NS_LOG_FUNCTION (this);
	switch (m_state)
	{
	case ASSOCIATED:
		return;
		break;
	case WAIT_PROBE_RESP:
		/* we have sent a probe request earlier so we
		 do not need to re-send a probe request immediately.
		 We just need to wait until probe-request-timeout
		 or until we get a probe response
		 */
		break;
	case BEACON_MISSED:
		/* we were associated but we missed a bunch of beacons
		 * so we should assume we are not associated anymore.
		 * We try to initiate a probe request now.
		 */
		m_linkDown();
		SetState(WAIT_PROBE_RESP);
		SendProbeRequest();
		break;
	case WAIT_ASSOC_RESP:
		/* we have sent an assoc request so we do not need to
		 re-send an assoc request right now. We just need to
		 wait until either assoc-request-timeout or until
		 we get an assoc response.
		 */
		break;
	case REFUSED:
		/* we have sent an assoc request and received a negative
		 assoc resp. We wait until someone restarts an
		 association with a given ssid.
		 */
		break;
	}
}

void StaCtrlWifiMac::AssocRequestTimeout(void)
{
	NS_LOG_FUNCTION (this);
	SetState(WAIT_ASSOC_RESP);
	SendAssociationRequest();
}

void StaCtrlWifiMac::ProbeRequestTimeout(void)
{
	NS_LOG_FUNCTION (this);
	SetState(WAIT_PROBE_RESP);
	SendProbeRequest();
}

void StaCtrlWifiMac::MissedBeacons(void)
{
	NS_LOG_FUNCTION (this);
	if (m_beaconWatchdogEnd > Simulator::Now())
	{
		m_beaconWatchdog = Simulator::Schedule(
				m_beaconWatchdogEnd - Simulator::Now(),
				&StaCtrlWifiMac::MissedBeacons, this);
		return;
	} NS_LOG_DEBUG ("beacon missed");
	SetState(BEACON_MISSED);
	TryToEnsureAssociated();
}

void StaCtrlWifiMac::RestartBeaconWatchdog(Time delay)
{
	NS_LOG_FUNCTION (this << delay);
	m_beaconWatchdogEnd = std::max(Simulator::Now() + delay,
			m_beaconWatchdogEnd);
	if (Simulator::GetDelayLeft(m_beaconWatchdog) < delay
			&& m_beaconWatchdog.IsExpired())
	{
		NS_LOG_DEBUG ("really restart watchdog.");
		m_beaconWatchdog = Simulator::Schedule(delay,
				&StaCtrlWifiMac::MissedBeacons, this);
	}
}

bool StaCtrlWifiMac::IsAssociated(void) const
{
	return m_state == ASSOCIATED;
}

bool StaCtrlWifiMac::IsWaitAssocResp(void) const
{
	return m_state == WAIT_ASSOC_RESP;
}

void StaCtrlWifiMac::Enqueue(Ptr<const Packet> packet, Mac48Address to)
{
	NS_LOG_FUNCTION (this << packet << to);
	if (!IsAssociated())
	{
		NotifyTxDrop(packet);
		TryToEnsureAssociated();
		return;
	}
	WifiMacHeader hdr;

	// If we are not a QoS AP then we definitely want to use AC_BE to
	// transmit the packet. A TID of zero will map to AC_BE (through \c
	// QosUtilsMapTidToAc()), so we use that as our default here.
	uint8_t tid = 0;

	// For now, an AP that supports QoS does not support non-QoS
	// associations, and vice versa. In future the AP model should
	// support simultaneously associated QoS and non-QoS STAs, at which
	// point there will need to be per-association QoS state maintained
	// by the association state machine, and consulted here.
	if (m_qosSupported)
	{
		hdr.SetType(WIFI_MAC_QOSDATA);
		hdr.SetQosAckPolicy(WifiMacHeader::NORMAL_ACK);
		hdr.SetQosNoEosp();
		hdr.SetQosNoAmsdu();
		// Transmission of multiple frames in the same TXOP is not
		// supported for now
		hdr.SetQosTxopLimit(0);

		// Fill in the QoS control field in the MAC header
		tid = QosUtilsGetTidForPacket(packet);
		// Any value greater than 7 is invalid and likely indicates that
		// the packet had no QoS tag, so we revert to zero, which'll
		// mean that AC_BE is used.
		if (tid >= 7)
		{
			tid = 0;
		}
		hdr.SetQosTid(tid);
	}
	else
	{
		hdr.SetTypeData();
	}

	hdr.SetAddr1(GetBssid());
	hdr.SetAddr2(m_low->GetAddress());
	hdr.SetAddr3(to);
	hdr.SetDsNotFrom();
	hdr.SetDsTo();

	if (m_qosSupported)
	{
		// Sanity check that the TID is valid
		NS_ASSERT(tid < 8);
		m_edca[QosUtilsMapTidToAc(tid)]->Queue(packet, hdr);
	}
	else
	{
		m_dca->Queue(packet, hdr);
	}
}

void StaCtrlWifiMac::Receive(Ptr<Packet> packet, const WifiMacHeader *hdr)
{
	NS_LOG_FUNCTION (this << packet << hdr);

	// normal process
	NS_ASSERT(!hdr->IsCtl ());
	if (hdr->GetAddr3() == GetAddress())
	{
		NS_LOG_LOGIC ("packet sent by us.");
		return;
	}
	else if (hdr->GetAddr1() != GetAddress() && !hdr->GetAddr1().IsGroup())
	{
		NS_LOG_LOGIC ("packet is not for us");
		NotifyRxDrop(packet);
		return;
	}
	else if (hdr->IsData())
	{
		// special process
		WifiMacHeader myHdr;
		packet->PeekHeader(myHdr);
		if (myHdr.IsRts() || myHdr.IsCts()) //Rts for data channel
		{
			m_dca->NotifyDataChannel(packet, *hdr);
			return; // do NOT go to the normal process
		}

		if (!IsAssociated())
		{
			NS_LOG_LOGIC ("Received data frame while not associated: ignore");
			NotifyRxDrop(packet);
			return;
		}
		if (!(hdr->IsFromDs() && !hdr->IsToDs()))
		{
			NS_LOG_LOGIC ("Received data frame not from the DS: ignore");
			NotifyRxDrop(packet);
			return;
		}
		if (hdr->GetAddr2() != GetBssid())
		{
			NS_LOG_LOGIC ("Received data frame not from the BSS we are associated with: ignore");
			NotifyRxDrop(packet);
			return;
		}

		if (hdr->IsQosData())
		{
			if (hdr->IsQosAmsdu())
			{
				NS_ASSERT(hdr->GetAddr3 () == GetBssid ());
				DeaggregateAmsduAndForward(packet, hdr);
				packet = 0;
			}
			else
			{
				ForwardUp(packet, hdr->GetAddr3(), hdr->GetAddr1());
			}
		}
		else
		{
			ForwardUp(packet, hdr->GetAddr3(), hdr->GetAddr1());
		}
		return;
	}
	else if (hdr->IsProbeReq() || hdr->IsAssocReq())
	{
		// This is a frame aimed at an AP, so we can safely ignore it.
		NotifyRxDrop(packet);
		return;
	}
	else if (hdr->IsBeacon())
	{
		MgtBeaconHeader beacon;
		packet->RemoveHeader(beacon);
		bool goodBeacon = false;
		if (GetSsid().IsBroadcast() || beacon.GetSsid().IsEqual(GetSsid()))
		{
			goodBeacon = true;
		}
		if ((IsWaitAssocResp() || IsAssociated())
				&& hdr->GetAddr3() != GetBssid())
		{
			goodBeacon = false;
		}
		if (goodBeacon)
		{
			Time delay = MicroSeconds(
					beacon.GetBeaconIntervalUs() * m_maxMissedBeacons);
			RestartBeaconWatchdog(delay);
			SetBssid(hdr->GetAddr3());
		}
		if (goodBeacon && m_state == BEACON_MISSED)
		{
			SetState(WAIT_ASSOC_RESP);
			SendAssociationRequest();
		}
		return;
	}
	else if (hdr->IsProbeResp())
	{
		if (m_state == WAIT_PROBE_RESP)
		{
			MgtProbeResponseHeader probeResp;
			packet->RemoveHeader(probeResp);
			if (!probeResp.GetSsid().IsEqual(GetSsid()))
			{
				//not a probe resp for our ssid.
				return;
			}
			SetBssid(hdr->GetAddr3());
			Time delay = MicroSeconds(
					probeResp.GetBeaconIntervalUs() * m_maxMissedBeacons);
			RestartBeaconWatchdog(delay);
			if (m_probeRequestEvent.IsRunning())
			{
				m_probeRequestEvent.Cancel();
			}
			SetState(WAIT_ASSOC_RESP);
			SendAssociationRequest();
		}
		return;
	}
	else if (hdr->IsAssocResp())
	{
		if (m_state == WAIT_ASSOC_RESP)
		{
			MgtAssocResponseHeader assocResp;
			packet->RemoveHeader(assocResp);
			if (m_assocRequestEvent.IsRunning())
			{
				m_assocRequestEvent.Cancel();
			}
			if (assocResp.GetStatusCode().IsSuccess())
			{
				SetState(ASSOCIATED);
				NS_LOG_DEBUG ("assoc completed");
				SupportedRates rates = assocResp.GetSupportedRates();
				for (uint32_t i = 0; i < m_phy->GetNModes(); i++)
				{
					WifiMode mode = m_phy->GetMode(i);
					if (rates.IsSupportedRate(mode.GetDataRate()))
					{
						m_stationManager->AddSupportedMode(hdr->GetAddr2(),
								mode);
						if (rates.IsBasicRate(mode.GetDataRate()))
						{
							m_stationManager->AddBasicMode(mode);
						}
					}
				}
				if (!m_linkUp.IsNull())
				{
					m_linkUp();
				}
			}
			else
			{
				NS_LOG_DEBUG ("assoc refused");
				SetState(REFUSED);
			}
		}
		return;
	}

	// Invoke the receive handler of our parent class to deal with any
	// other frames. Specifically, this will handle Block Ack-related
	// Management Action frames.
	RegularCtrlWifiMac::Receive(packet, hdr);
}

SupportedRates StaCtrlWifiMac::GetSupportedRates(void) const
{
	SupportedRates rates;
	for (uint32_t i = 0; i < m_phy->GetNModes(); i++)
	{
		WifiMode mode = m_phy->GetMode(i);
		rates.AddSupportedRate(mode.GetDataRate());
	}
	return rates;
}

void StaCtrlWifiMac::SetState(MacState value)
{
	if (value == ASSOCIATED && m_state != ASSOCIATED)
	{
		m_assocLogger(GetBssid());
	}
	else if (value != ASSOCIATED && m_state == ASSOCIATED)
	{
		m_deAssocLogger(GetBssid());
	}
	m_state = value;
}
} // namespace ns3
